package com.gitee.wsl.flow.transform

import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch

private val NULL = Any()

fun <T> Flow<T>.rest(time: Long, emitLastEmissionWithoutRest: Boolean = false): Flow<T> {
    return channelFlow {
        var upStreamFinished = false
        var lastValue: Any? = NULL
        suspend fun pushValue() {
            val value = lastValue
            if (lastValue !== NULL) {
                lastValue = NULL
                send(value as T)
            }
        }

        val ticker = launch {
            while (isActive) {
                delay(time)
                pushValue()
                if (upStreamFinished) {
                    break
                }
            }
            close()
        }
        launch {
            collect {
                lastValue = it
            }
            if (emitLastEmissionWithoutRest) {
                pushValue()
                ticker.cancel()
            }
            upStreamFinished = true
        }
    }
}